from openmind_hub import RepositoryNotFoundError, GatedRepoError
from openmind_hub import create_repo, delete_repo, whoami, upload_file, upload_folder, get_om_file_metadata, \
    CommitOperationAdd, create_commit, list_repo_tree
import unittest
import time
import os,string,random

token = '123456789()'
repo_id = "guoxiaozhen2/dataset_test2024"
license = "mit"

def generate_large_file(filename, size_in_bytes):
    with open(filename, 'wb') as file:
        chunk_size = 1024 * 1024  # 1MB
        chunks = size_in_bytes // chunk_size
        remainder = size_in_bytes % chunk_size
        data = b'\0' * chunk_size

        for _ in range(chunks):
            file.write(data)

        if remainder > 0:
            file.write(data[:remainder])

class TestOpenmindHub(unittest.TestCase):
    def test_01_whoami_success(self):
        try:
            a=whoami(
                token=token,
            )
            self.assertEqual(a["username"],"guoxiaozhen2","用户查验有误")
        except:
            raise Exception("查询用户信息报错")
    def test_02_whoami_wrongtoken(self):
        try:
            whoami(
                token=token+"A",
            )
            raise Exception("查询用户信息未校验token")
        except:
            pass
    def test_03_uploadfile_success(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/requirements.txt",
                path_in_repo="a/requirements.txt",
                repo_id=repo_id,
                token=token,
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
            time.sleep(3)
            a = get_om_file_metadata(
                url="https://openmind.test.osinfra.cn/api/v1/file/{}/info?ref=main&path=a/requirements.txt".format(repo_id),
                token=token,
                proxies=None,
                timeout=5,
            )
            self.assertIsNotNone(a,"文件不存在，上传失败")
        except :
            raise Exception("上传文件报错")
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
    def test_04_uploadfile_noexistropo(self):
        try:
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/requirements.txt",
                path_in_repo="a/requirements.txt",
                repo_id=repo_id+"A",
                token=token,
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
            raise Exception("上传文件到不存在的仓库未报错")
        except RepositoryNotFoundError:
            pass
    def test_05_uploadfile_wrongtoken(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(2)
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/requirements.txt",
                path_in_repo="a/requirements.txt",
                repo_id=repo_id,
                token=token + "A",
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
            raise Exception("上传文件未校验token")
        except GatedRepoError:
            pass
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
    def test_06_get_om_file_metadata_wrongurl(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/requirements.txt",
                path_in_repo="a/requirements.txt",
                repo_id=repo_id,
                token=token,
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
            time.sleep(3)
            get_om_file_metadata(
                url="https://openmind.test.osinfra.cn/api/v1/file/{}/info?ref=main&path=a/requirements.txt".format(repo_id+"A"),
                token=token,
                proxies=None,
                timeout=5,
            )
            raise Exception("获取文件错误的URL未报错")
        except RepositoryNotFoundError:
            pass
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
    def test_07_get_om_file_metadata_wrongtoken(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/requirements.txt",
                path_in_repo="a/requirements.txt",
                repo_id=repo_id,
                token=token,
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
            time.sleep(3)
            a = get_om_file_metadata(
                url="https://telecom.openmind.cn/api/v1/file/{}/info?ref=main&path=a/requirements.txt".format(repo_id),
                token=token+"A",
                proxies=None,
                timeout=5,
            )
            raise Exception("获取文件错误的token未报错")
        except GatedRepoError:
            pass
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
    def test_08_create_commit_success(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            operations = [CommitOperationAdd(path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/requirements.txt", path_in_repo="test1.txt")]
            create_commit(
                repo_id=repo_id,
                operations=operations,
                commit_message="create commit",
                commit_description=None,
                token=token,
                revision="main",
                num_threads=5,
            )
            time.sleep(3)
            a = get_om_file_metadata(
                url="https://openmind.test.osinfra.cn/api/v1/file/{}/info?ref=main&path=test1.txt".format(repo_id),
                token=token,
                proxies=None,
                timeout=5,
            )
            self.assertIsNotNone(a, "文件不存在，提交失败")
        except:
            raise Exception("上传文件报错")
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
    def test_09_create_commit_wrongtoken(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            operations = [CommitOperationAdd(path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/requirements.txt", path_in_repo="test1.txt")]
            create_commit(
                repo_id=repo_id,
                operations=operations,
                commit_message="create commit",
                commit_description=None,
                token=token+"a",
                revision="main",
                num_threads=5,
            )
            raise Exception("创建提交，错误的token未报错")
        except GatedRepoError:
            pass
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
    def test_10_create_commit_noexistrepo(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            operations = [CommitOperationAdd(path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/requirements.txt", path_in_repo="test1.txt")]
            create_commit(
                repo_id=repo_id+"a",
                operations=operations,
                commit_message="create commit",
                commit_description=None,
                token=token,
                revision="main",
                num_threads=5,
            )
            raise Exception("创建提交，不存在的仓库未报错")
        except RepositoryNotFoundError:
            pass
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
    def test_11_list_repo_tree_success(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=True,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/requirements.txt",
                path_in_repo="requirements.txt",
                repo_id=repo_id,
                token=token,
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
            time.sleep(3)
            a = list_repo_tree(
                    repo_id=repo_id,
                    path_in_repo=None,
                    revision=None,
                    token=token,
                )
            list_a=list(a)
            title = []
            for i in range(len(list_a)):
                title.append(list_a[i].path)
            self.assertIn("requirements.txt",title,"文件不存在，上传失败")
        except:
            raise Exception("获取文件树信息报错")
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
    def test_12_list_repo_tree_notoken(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=True,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/requirements.txt",
                path_in_repo="requirements.txt",
                repo_id=repo_id,
                token=token,
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
            time.sleep(3)
            list_repo_tree(
                    repo_id=repo_id,
                    path_in_repo=None,
                    revision=None,
                    token=None,
                )
        except GatedRepoError:
            pass
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
    def test_13_list_repo_tree_noexistrepo(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=True,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/requirements.txt",
                path_in_repo="requirements.txt",
                repo_id=repo_id,
                token=token,
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
            time.sleep(3)
            list_repo_tree(
                    repo_id=repo_id+"A",
                    path_in_repo=None,
                    revision=None,
                    token=None,
                )
        except RepositoryNotFoundError:
            pass
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
    def test_14_uploadfile_jsonl(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/requirements.jsonl",
                path_in_repo="a/requirements.jsonl",
                repo_id=repo_id,
                token=token,
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
            time.sleep(3)
            a = get_om_file_metadata(
                url="https://openmind.test.osinfra.cn/api/v1/file/{}/info?ref=main&path=a/requirements.jsonl".format(repo_id),
                token=token,
                proxies=None,
                timeout=5,
            )
            self.assertIsNotNone(a,"文件不存在，上传失败")
        except :
            raise Exception("上传文件报错")
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
    def test_15_uploadfile_duplicate(self):
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(1)
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/requirements.txt",
                path_in_repo="a/requirements.jsonl",
                repo_id=repo_id,
                token=token,
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
            time.sleep(5)
            upload_file(
                path_or_fileobj=r"/openmind-hub/auto_test/autotest/test_openmind_hub/requirements.txt",
                path_in_repo="a/requirements.jsonl",
                repo_id=repo_id,
                token=token,
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
        except :
            raise Exception("重复上传文件报错")
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
    def test_16_uploadfile_largefile_5G(self):
        file_path = "/openmind-hub/auto_test/autotest/test_openmind_hub/filedata/testfile"
        size_in_bytes = int(5*1024*1024*1024)
        try:
            generate_large_file(file_path, size_in_bytes)
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            upload_file(
                path_or_fileobj=file_path,
                path_in_repo="a/filedata.txt",
                repo_id=repo_id,
                token=token,
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
            time.sleep(3)
            a = get_om_file_metadata(
                url="https://openmind.test.osinfra.cn/api/v1/file/{}/info?ref=main&path=a/filedata.txt".format(repo_id),
                token=token,
                proxies=None,
                timeout=5,
            )
            self.assertIsNotNone(a,"文件不存在，上传失败")
        except :
            raise Exception("上传文件报错")
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
            os.remove(file_path)
    def test_17_uploadfile_largefile_50G(self):
        file_path = "/openmind-hub/auto_test/autotest/test_openmind_hub/filedata/testfile"
        size_in_bytes = int(50*1024*1024*1024)
        try:
            generate_large_file(file_path, size_in_bytes)
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            upload_file(
                path_or_fileobj=file_path,
                path_in_repo="a/filedata.txt",
                repo_id=repo_id,
                token=token,
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
            time.sleep(3)
            a = get_om_file_metadata(
                url="https://openmind.test.osinfra.cn/api/v1/file/{}/info?ref=main&path=a/filedata.txt".format(repo_id),
                token=token,
                proxies=None,
                timeout=5,
            )
            self.assertIsNotNone(a,"文件不存在，上传失败")
        except :
            raise Exception("上传文件报错")
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
            os.remove(file_path)
    def test_18_uploadfile_largefile_5Gmulti10(self):
        file_path = []
        for i in range(5):
            file_path.append("/openmind-hub/auto_test/autotest/test_openmind_hub/filedata/test{}".format(i))
        size_in_bytes = int(5*1024*1024*1024)
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            for i in file_path:
                generate_large_file(i, size_in_bytes)
            upload_folder(
                repo_id=repo_id,
                folder_path="/openmind-hub/auto_test/autotest/test_openmind_hub/filedata",
                path_in_repo="/filedata",
                commit_message="upload folder",
                commit_description=None,
                token=token,
                revision=None,
                # allow_patterns=["*.bin", "*.py"],
                ignore_patterns="*.log",
                num_threads=5,
            )
        except :
            raise Exception("上传文件报错")
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
            for i in file_path:
                os.remove(i)
    def test_19_upload_folder_success(self):
        file_path = []
        for i in range(10):
            file_path.append("/openmind-hub/auto_test/autotest/test_openmind_hub/filedata/test{}".format(i))
        size_in_bytes = int(0.00001 * 1024 * 1024 * 1024)
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            for i in file_path:
                generate_large_file(i, size_in_bytes)
            upload_folder(
                repo_id=repo_id,
                folder_path="/openmind-hub/auto_test/autotest/test_openmind_hub/filedata",
                path_in_repo="/filedata",
                commit_message="upload folder",
                commit_description=None,
                token=token,
                revision=None,
                # allow_patterns=["*.bin", "*.py"],
                ignore_patterns="*.log",
                num_threads=5,
            )
        except:
            raise Exception("上传文件报错")
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
            for i in file_path:
                os.remove(i)
    def test_20_upload_folder_multi2000(self):
        file_dir = []
        # 创建4000个随机文件
        for i in range(2000):
            # 生成随机文件名
            random_file_name = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
            file_path = f"/openmind-hub/auto_test/autotest/test_openmind_hub/filedata/{random_file_name}.txt"
            file_dir.append(file_path)
            # 创建空白文件
            with open(file_path, 'w'):
                pass
        print("已创建2000个随机文件")
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            upload_folder(
                repo_id=repo_id,
                folder_path="/openmind-hub/auto_test/autotest/test_openmind_hub/filedata",
                path_in_repo="/filedata",
                commit_message="upload folder",
                commit_description=None,
                token=token,
                revision=None,
                # allow_patterns=["*.bin", "*.py"],
                ignore_patterns="*.log",
                num_threads=5,
            )
        except:
            raise Exception("上传文件报错")
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
            for i in file_dir:
                os.remove(i)
if __name__ == '__main__':
    unittest.main()